#!/usr/bin/python
# Filename: tag.py

import math;
import sys;

def print_tree(tree):
	if len(tree) == 2:
		return '[\"'+tree[0]+'\", \"'+tree[1]+'\"]';
	else:
		return '[\"'+tree[0]+'\", '+print_tree(tree[1])+', '+print_tree(tree[2])+']';

def max_dict(s):
	value = float('-inf');
	key = '';
	for name in s:
		if s[name] >= value:
			key = name;
			value = s[name];
	return (key,value);

def CKY(x,un,bi,noterm,words): #the CKY algorithm
	pi = {};
	bp = {};
	pi_dict = {};
	n=len(x);
	
	for i in range(1,len(x)+1): #initialization
		bp[i,i]={};
		pi[i,i]={}
		if not(x[i] in words):  #replace the rare words with '_RARE_'
			w_rep = '_RARE_';
		else:
			w_rep = x[i];
		for term in noterm:
			if not (term in un):
				pi[i,i][term]=float('-inf');
				bp[i,i][term]=[];
			else:
				if not (w_rep in un[term]):#if X->xi does not exist
					pi[i,i][term] = float('-inf');
					bp[i,i][term] = [];
				else:
					pi[i,i][term] = un[term][w_rep];
					bp[i,i][term] = [term, x[i]];


	for l in range(1,n): #handle with pi and bp
		for i in range(1,n-l+1):
			j=i+l;
			bp[i,j]={};
			pi[i,j]={};
			for x_ntm in noterm:
				if not(x_ntm in bi):
					pi[i,j][x_ntm]=float('-inf');
					continue;
				value = float('-inf');
				for terms in bi[x_ntm]:
					y_ntm = terms[0];
					z_ntm = terms[1];
					for s in range(i,j):
						temp = bi[x_ntm][y_ntm,z_ntm]+pi[i,s][y_ntm]+pi[s+1,j][z_ntm];
						if temp > value:
							value = temp;
							bp[i,j][x_ntm] = [x_ntm, bp[i,s][y_ntm], bp[s+1,j][z_ntm]];
				pi[i,j][x_ntm] = value;
	
	if pi[1,n]['S'] != float('-inf'):
		return bp[1,n]['S'];
	else:
		(k,v) = max_dict(pi[1,n]);
		return bp[1,n][k];
		

un_tbl = dict();
bi_tbl = dict();
no_set = set();
word_set = set();

fr_un_dat = file(sys.argv[1]);
for line in fr_un_dat:    #read into un_tbl and add all the nonterminals into set
	words = line.strip().split(' ');
	if words[0] in un_tbl:
		un_tbl[words[0]][words[1]] = math.log(float(words[2]));
	else:
		un_tbl[words[0]] = {words[1] : math.log(float(words[2]))}; 
	if not(no_set.__contains__(words[0])):
			no_set.add(words[0]);
	if not(word_set.__contains__(words[1])):
			word_set.add(words[1]);
fr_un_dat.close();


fr_bi_dat = file(sys.argv[2]);
for line in fr_bi_dat: 	#read into bi_tbl and add nonterminals into set
	words = line.strip().split(' ');
	if words[0] in bi_tbl:
		bi_tbl[words[0]][words[1],words[2]] = math.log(float(words[3]));
	else:
		bi_tbl[words[0]]={(words[1],words[2]):math.log(float(words[3]))};
	for j in range(0,3):
		if not(no_set.__contains__(words[j])):
			no_set.add(words[j]);
fr_bi_dat.close();



fr_dev = file(sys.argv[3]);

fw_dev = open(sys.argv[4],'w');

sentence = {};
i = 1; 
for line in fr_dev:
	words = line.strip().split(' ');
	for word in words: #read in the sentence and start index from 1
		sentence[i] = word;
		i += 1;
	#use the sentence as an input of CKY Algorithm
	x = sentence.copy();
	tree = CKY(x,un_tbl,bi_tbl,no_set,word_set);
	line = print_tree(tree);
	fw_dev.write(line);
	i = 1;
	fw_dev.write('\n');
	sentence.clear();
fw_dev.close();
fr_dev.close();
		
